from osgeo import gdal
from sklearn.linear_model import LinearRegression
from scipy import stats
import numpy as np
import os
from tqdm import tqdm

'''
本代码全部由deepseek v3生成：
代码简介
1. 代码功能
该代码用于对时间序列的NDVI（归一化植被指数）数据进行多元回归分析，预测未来的NDVI值，并评估回归模型的性能。具体功能包括：
从指定目录读取NDVI、温度、降水和蒸散发数据。
对每个像元进行多元回归分析，计算回归系数、截距、R²（决定系数）和t检验的t值。
将回归系数、截距、R²和t值保存为GeoTIFF图像。
使用回归模型预测未来的NDVI值，并将结果保存为GeoTIFF图像。
提供进度条显示，方便用户了解程序运行进度。

2. 代码结构
代码由多个函数组成，每个函数负责一个独立的功能模块：
read_data_from_directory: 从指定目录读取数据文件，并自动处理NoData值和inf值。
read_data: 从指定目录读取特定年份的数据。
perform_regression_with_stats: 执行多元回归分析，计算回归系数、截距、R²和t值。
save_image: 将数据保存为GeoTIFF图像，并保留地理信息和NoData值。
predict_ndvi: 使用回归模型预测NDVI值，并将结果裁剪到[0, 1]范围内。
clip_ndvi: 将NDVI值裁剪到[0, 1]范围内。
get_template_file: 获取指定目录下的第一个.tif文件作为模板文件。
main: 主函数，负责调用上述函数完成整个流程。

3. 输入数据
NDVI数据: 年均值图像，时间长度为2000-2020年。
温度数据: 年均温图像，时间长度为2000-2020年。
降水数据: 年总降水图像，时间长度为2000-2020年。
蒸散发数据: 年总蒸散发图像，时间长度为2000-2020年。

4. 输出结果
回归系数: 温度、降水和蒸散发对NDVI的影响系数，保存为temp_coefficient.tif、precip_coefficient.tif和et_coefficient.tif。
截距: 回归模型的截距，保存为intercept.tif。
R²值: 回归模型的决定系数，保存为r2_values.tif。
t值: t检验的t值，分别保存为t_values_temp.tif、t_values_precip.tif和t_values_et.tif。
预测的NDVI值: 预测的未来NDVI值，保存为ndvi_pred.tif。

5. 技术细节
数据处理: 使用GDAL库读取和保存GeoTIFF文件，自动处理NoData值和inf值。
回归分析: 使用scikit-learn库进行多元线性回归分析，计算回归系数、截距和R²。
t检验: 使用scipy库计算t检验的t值。
进度显示: 使用tqdm库显示逐像元回归和预测的进度。
结果裁剪: 将预测的NDVI值裁剪到[0, 1]范围内，确保结果符合NDVI的定义。

6. 适用场景
生态环境研究: 分析温度、降水和蒸散发对植被覆盖的影响。
农业监测: 预测未来植被覆盖情况，评估农业生产潜力。
气候变化研究: 研究气候变化对植被的影响。

7. 代码优势
模块化设计: 每个功能模块独立，易于维护和扩展。
自动化处理: 自动处理NoData值和inf值，减少人工干预。
高效计算: 使用逐像元回归分析，支持大规模数据处理。
可视化输出: 将结果保存为GeoTIFF图像，方便在GIS软件中可视化。

8. 运行环境
Python版本: 3.7及以上。
依赖库: numpy, gdal, scipy, scikit-learn, tqdm。
通过该代码，用户可以高效地完成NDVI数据的回归分析和预测，并获得丰富的统计结果和可视化输出。
'''


#1.定义函数
#1.1读取数据的函数
def read_data_from_directory(directory, keyword=None, extension=".tif"):
    """
    从指定目录读取数据文件，并自动获取NoData值替换为NaN
    :param directory: 数据文件所在的目录
    :param keyword: 文件名中的关键字（可选，用于筛选特定文件）
    :param extension: 文件扩展名（默认读取.tif文件）
    :return: 数据列表（每个元素是一个NumPy数组）
    """
    data_list = []
    file_paths = []

    # 遍历目录下的所有文件
    for file_name in os.listdir(directory):
        # 检查文件扩展名和关键字
        if file_name.endswith(extension) and (keyword is None or keyword in file_name):
            file_path = os.path.join(directory, file_name)
            file_paths.append(file_path)

    # 按文件名排序，确保数据按时间顺序排列
    file_paths.sort()

    # 读取数据
    for file_path in file_paths:
        dataset = gdal.Open(file_path)
        if dataset is None:
            raise FileNotFoundError(f"文件 {file_path} 无法读取")

        band = dataset.GetRasterBand(1)
        data = band.ReadAsArray()

        # 获取NoData值
        no_data_value = band.GetNoDataValue()

        # 如果存在NoData值，将其替换为NaN
        if no_data_value is not None:
            data[data == no_data_value] = np.nan

        # 检查是否存在inf值，并将其替换为NaN
        if np.isinf(data).any():
            data[np.isinf(data)] = np.nan

        data_list.append(data)

    return data_list


#1.2多元回归分析的函数，该函数在原有回归分析的基础上，计算R²和t检验的t值，并保存为图像。
def perform_regression_with_stats(ndvi_data, temp_data, precip_data, et_data):
    """
    执行多元回归分析，并计算R²和t检验的t值
    :param ndvi_data: NDVI数据（三维数组，形状为(years, rows, cols)）
    :param temp_data: 温度数据（三维数组，形状为(years, rows, cols)）
    :param precip_data: 降水数据（三维数组，形状为(years, rows, cols)）
    :param et_data: 蒸散发数据（三维数组，形状为(years, rows, cols)）
    :return: 回归系数（形状为(rows, cols, 3)），截距（形状为(rows, cols)），R²（形状为(rows, cols)），t值（形状为(rows, cols, 3)）
    """
    #from sklearn.linear_model import LinearRegression

    # 获取图像的形状
    years, rows, cols = ndvi_data.shape

    # 将数据展平为二维数组 (years, rows*cols)
    ndvi_flat = ndvi_data.reshape(years, -1)
    temp_flat = temp_data.reshape(years, -1)
    precip_flat = precip_data.reshape(years, -1)
    et_flat = et_data.reshape(years, -1)

    # 初始化回归模型
    model = LinearRegression()

    # 逐像元回归
    coefficients = np.zeros((rows * cols, 3))  # 存储回归系数
    intercepts = np.zeros(rows * cols)  # 存储截距
    r2_values = np.zeros(rows * cols)  # 存储R²
    t_values = np.zeros((rows * cols, 3))  # 存储t值

    # 使用tqdm显示进度
    for i in tqdm(range(rows * cols), desc="逐像元回归", unit="像元"):
        # 提取当前像元的自变量和因变量
        X = np.column_stack((temp_flat[:, i], precip_flat[:, i], et_flat[:, i]))
        y = ndvi_flat[:, i]

        # 检查是否存在NaN或inf值
        valid = ~np.isnan(X).any(axis=1) & ~np.isnan(y) & ~np.isinf(X).any(axis=1) & ~np.isinf(y)
        X_valid = X[valid]
        y_valid = y[valid]

        # 如果有有效数据，进行回归
        if len(X_valid) > 0:
            model.fit(X_valid, y_valid)
            coefficients[i] = model.coef_
            intercepts[i] = model.intercept_

            # 计算R²
            y_pred = model.predict(X_valid)
            r2_values[i] = model.score(X_valid, y_valid)

            # 计算t值
            n = len(X_valid)
            p = X_valid.shape[1] + 1  # 包括截距
            residuals = y_valid - y_pred
            mse = np.sum(residuals ** 2) / (n - p)
            cov = np.linalg.inv(np.dot(X_valid.T, X_valid))
            se = np.sqrt(np.diag(mse * cov))
            t_values[i] = model.coef_ / se
        else:
            # 如果没有有效数据，将系数、截距、R²和t值设置为NaN
            coefficients[i] = np.nan
            intercepts[i] = np.nan
            r2_values[i] = np.nan
            t_values[i] = np.nan

    # 将结果重新整形为图像形状
    coefficients = coefficients.reshape(rows, cols, 3)
    intercepts = intercepts.reshape(rows, cols)
    r2_values = r2_values.reshape(rows, cols)
    t_values = t_values.reshape(rows, cols, 3)

    return coefficients, intercepts, r2_values, t_values


#1.3 保存结果的函数
def save_image(data, file_name, template_file, output_dir="output"):
    """
    保存图像数据，并使用模板文件的NoData值
    :param data: 要保存的数据（NumPy数组）
    :param file_name: 保存的文件名
    :param template_file: 模板文件路径（用于获取NoData值和地理信息）
    :param output_dir: 输出目录（默认为"output"）
    """
    # 创建输出目录（如果不存在）
    os.makedirs(output_dir, exist_ok=True)

    # 构建完整的文件路径
    file_path = os.path.join(output_dir, file_name)

    # 打开模板文件
    template_dataset = gdal.Open(template_file)
    if template_dataset is None:
        raise FileNotFoundError(f"模板文件 {template_file} 未找到")

    # 获取模板文件的NoData值
    band = template_dataset.GetRasterBand(1)
    no_data_value = band.GetNoDataValue()

    # 创建输出文件
    driver = gdal.GetDriverByName('GTiff')
    dataset = driver.Create(file_path, data.shape[1], data.shape[0], 1, gdal.GDT_Float32)

    # 设置地理信息和投影
    dataset.SetGeoTransform(template_dataset.GetGeoTransform())
    dataset.SetProjection(template_dataset.GetProjection())

    # 写入数据
    band = dataset.GetRasterBand(1)
    data[np.isnan(data)] = no_data_value  # 将NaN替换为NoData值
    band.WriteArray(data)

    # 设置NoData值
    band.SetNoDataValue(no_data_value)
    dataset.FlushCache()

#1.4 在获取template_file时，确保只选择.tif文件
def get_template_file(directory, extension=".tif"):
    """
    获取指定目录下的第一个.tif文件作为模板文件
    :param directory: 目录路径
    :param extension: 文件扩展名（默认为.tif）
    :return: 模板文件路径
    """
    # 获取目录下所有.tif文件
    tif_files = [f for f in os.listdir(directory) if f.endswith(extension)]

    if not tif_files:
        raise FileNotFoundError(f"目录 {directory} 中没有找到{extension}文件")

    # 返回第一个.tif文件的完整路径
    return os.path.join(directory, tif_files[0])

#1.5 预测NDVI的函数
def predict_ndvi(coefficients, intercepts, temp_data, precip_data, et_data, template_file, output_dir="output"):
    """
    预测NDVI
    :param coefficients: 回归系数（形状为(rows, cols, 3)）
    :param intercepts: 截距（形状为(rows, cols)）
    :param temp_data: 温度数据（二维数组，形状为(rows, cols)）
    :param precip_data: 降水数据（二维数组，形状为(rows, cols)）
    :param et_data: 蒸散发数据（二维数组，形状为(rows, cols)）
    :param template_file: 模板文件路径（用于保存结果）
    :param output_dir: 输出目录（默认为"output"）
    :return: 预测的NDVI数据（二维数组，形状为(rows, cols)）
    """
    # 展平数据
    temp_flat = temp_data.reshape(-1)
    precip_flat = precip_data.reshape(-1)
    et_flat = et_data.reshape(-1)

    # 预测NDVI
    rows, cols = intercepts.shape
    ndvi_pred = np.zeros(rows * cols)

    # 使用tqdm显示进度
    for i in tqdm(range(rows * cols), desc="预测NDVI", unit="像元"):
        # 检查是否存在NaN或inf值
        if not np.isnan(coefficients.reshape(-1, 3)[i]).any() and \
           not np.isnan(temp_flat[i]) and \
           not np.isnan(precip_flat[i]) and \
           not np.isnan(et_flat[i]) and \
           not np.isinf(temp_flat[i]) and \
           not np.isinf(precip_flat[i]) and \
           not np.isinf(et_flat[i]):
            ndvi_pred[i] = (coefficients.reshape(-1, 3)[i][0] * temp_flat[i] +
                            coefficients.reshape(-1, 3)[i][1] * precip_flat[i] +
                            coefficients.reshape(-1, 3)[i][2] * et_flat[i] +
                            intercepts.reshape(-1)[i])
        else:
            ndvi_pred[i] = np.nan

    # 将预测结果重新整形为图像形状
    ndvi_pred = ndvi_pred.reshape(rows, cols)

    # 裁剪NDVI值到[0, 1]范围
    ndvi_pred = clip_ndvi(ndvi_pred)

    # 保存预测结果
    save_image(ndvi_pred, "ndvi_pred.tif", template_file, output_dir)

    return ndvi_pred

#1.6 定义clip_ndvi函数，将NDVI值裁剪到[0, 1]范围内
def clip_ndvi(ndvi_data):
    """
    将NDVI值裁剪到[0, 1]范围内
    :param ndvi_data: 输入的NDVI数据（NumPy数组）
    :return: 裁剪后的NDVI数据（NumPy数组）
    """
    # 将大于1的值设置为1
    ndvi_data[ndvi_data > 1] = 1
    # 将小于0的值设置为0
    ndvi_data[ndvi_data < 0] = 0
    return ndvi_data

# 2.主程序,在主程序中调用上述函数完成整个流程。
if __name__ == '__main__':
    # 定义数据目录
    ndvi_dir = r"F:\test\NDVI2"
    precip_dir = r'F:\test\pre'
    temp_dir = r'F:\test\tmp'
    et_dir = r'F:\test\pet'
    predict_dir = r'F:\test\ssp585'
    output_dir = r'F:\test\ssp585_pred'  # 输出目录

    # 用于获取情景的关键字
    keyword_tmp = "tmp"
    keyword_pre = "pre"
    keyword_pet = "pet"


    # 检查输出目录是否存在，如果不存在则创建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 读取NDVI数据
    ndvi_data = read_data_from_directory(ndvi_dir)
    ndvi_stack = np.stack(ndvi_data, axis=0)

    # 读取年均温数据
    temp_data = read_data_from_directory(temp_dir)
    temp_stack = np.stack(temp_data, axis=0)

    # 读取年总降水数据
    precip_data = read_data_from_directory(precip_dir)
    precip_stack = np.stack(precip_data, axis=0)

    # 读取年总蒸散发数据
    et_data = read_data_from_directory(et_dir)
    et_stack = np.stack(et_data, axis=0)

    # 执行多元回归分析并计算R²和t值
    # 执行多元回归分析并计算R²和t值
    coefficients, intercepts, r2_values, t_values = perform_regression_with_stats(
        ndvi_stack, temp_stack, precip_stack, et_stack)

    # 获取模板文件路径
    template_file = get_template_file(ndvi_dir)

    # 保存回归系数和截距
    save_image(coefficients[:, :, 0], "temp_coefficient.tif", template_file, output_dir)
    save_image(coefficients[:, :, 1], "precip_coefficient.tif", template_file, output_dir)
    save_image(coefficients[:, :, 2], "et_coefficient.tif", template_file, output_dir)
    save_image(intercepts, "intercept.tif", template_file, output_dir)

    # 保存R²和t值
    save_image(r2_values, "r2_values.tif", template_file, output_dir)
    for i, var_name in enumerate(["temp", "precip", "et"]):
        save_image(t_values[:, :, i], f"t_values_{var_name}.tif", template_file, output_dir)

    # 预测NDVI
    temp_pred = read_data_from_directory(predict_dir, keyword=keyword_tmp)[0]
    precip_pred = read_data_from_directory(predict_dir, keyword=keyword_pre)[0]
    et_pred = read_data_from_directory(predict_dir, keyword=keyword_pet)[0]

    ndvi_pred = predict_ndvi(coefficients, intercepts, temp_pred, precip_pred, et_pred, template_file, output_dir)

    print(f"NDVI预测完成，结果已保存到 {output_dir} 目录")

    """    
    # 预测2021-2060年的NDVI值
    for year in range(2021, 2061):
        print(f"正在预测 {year} 年的NDVI...")
        
        # 读取未来年份的温度、降水和蒸散发数据
        temp_future = read_data_from_directory(temp_dir, keyword=f"temp_{year}")
        precip_future = read_data_from_directory(precip_dir, keyword=f"precip_{year}")
        et_future = read_data_from_directory(et_dir, keyword=f"et_{year}")

        # 检查是否成功读取数据
        if len(temp_future) == 0 or len(precip_future) == 0 or len(et_future) == 0:
            print(f"警告：{year} 年的数据缺失，跳过该年份的预测。")
            continue

        # 提取数据（假设每年只有一个文件）
        temp_future = temp_future[0]
        precip_future = precip_future[0]
        et_future = et_future[0]

        # 预测NDVI
        ndvi_future_pred = predict_ndvi(coefficients, intercepts, temp_future, precip_future, et_future, template_file, output_dir)

        # 保存预测结果
        save_image(ndvi_future_pred, f"ndvi_pred_{year}.tif", template_file, output_dir)

    print(f"NDVI预测完成，结果已保存到 {output_dir} 目录")
    """
